import numpy as np
from collections import Counter
from typing import List, Dict, Any, Tuple
class Node:
    def __init__(self, feature=None, threshold=None, left=None, right=None, value=None):
        self.feature = feature  # 分裂特征的索引
        self.threshold = threshold  # 分裂阈值
        self.left = left  # 左子节点
        self.right = right  # 右子节点
        self.value = value  # 叶节点的值（预测的类别）
class DecisionTreeID3:
    def __init__(self, max_depth=None, min_samples_split=2):
        self.max_depth = max_depth
        self.min_samples_split = min_samples_split
        self.tree = None
        self.class_to_int = None  # 映射类别标签到整数
    def fit(self, X: np.ndarray, y: np.ndarray):
        # 将类别标签转换为整数
        unique_classes = np.unique(y)
        self.class_to_int = {cls: i for i, cls in enumerate(unique_classes)}
        y_int = np.array([self.class_to_int[cls] for cls in y])
        self.tree = self._build_tree(X, y_int)
    def _build_tree(self, X: np.ndarray, y: np.ndarray, depth=0):
        n_samples, n_features = X.shape
        n_labels = len(np.unique(y))
        # 停止条件
        if (self.max_depth is not None and depth >= self.max_depth) or n_labels == 1 or n_samples < self.min_samples_split:
            leaf_value = self._most_common_label(y)
            # 将整数值转换回类别标签
            leaf_value = list(self.class_to_int.keys())[list(self.class_to_int.values()).index(leaf_value)]
            return Node(value=leaf_value)
        # 计算信息增益，选择最佳特征
        feature_idxs = np.arange(n_features)
        best_feature, best_threshold = self._best_split(X, y, feature_idxs)
        # 没有特征可用于分裂，返回叶节点
        if best_feature is None:
            leaf_value = self._most_common_label(y)
            leaf_value = list(self.class_to_int.keys())[list(self.class_to_int.values()).index(leaf_value)]
            return Node(value=leaf_value)
        # 分裂数据
        left_idxs, right_idxs = self._split(X[:, best_feature], best_threshold)
        left = self._build_tree(X[left_idxs, :], y[left_idxs], depth + 1)
        right = self._build_tree(X[right_idxs, :], y[right_idxs], depth + 1)
        return Node(feature=best_feature, threshold=best_threshold, left=left, right=right)
    def _best_split(self, X: np.ndarray, y: np.ndarray, feature_idxs: np.ndarray) -> Tuple[int, float]:
        best_gain = -1
        split_idx, split_threshold = None, None
        for feature_idx in feature_idxs:
            X_column = X[:, feature_idx]
            thresholds = np.unique(X_column)
            for threshold in thresholds:
                gain = self._information_gain(y, X_column, threshold)

                if gain > best_gain:
                    best_gain = gain
                    split_idx = feature_idx
                    split_threshold = threshold
        return split_idx, split_threshold
    def _information_gain(self, y: np.ndarray, X_column: np.ndarray, threshold: float) -> float:
        # 计算父节点的熵
        parent_entropy = self._entropy(y)
        # 分裂数据
        left_idxs, right_idxs = self._split(X_column, threshold)
        if len(left_idxs) == 0 or len(right_idxs) == 0:
            return 0
        # 计算子节点的熵
        n = len(y)
        n_left, n_right = len(left_idxs), len(right_idxs)
        e_left, e_right = self._entropy(y[left_idxs]), self._entropy(y[right_idxs])
        child_entropy = (n_left / n) * e_left + (n_right / n) * e_right
        # 计算信息增益
        return parent_entropy - child_entropy
    def _split(self, X_column: np.ndarray, threshold: float) -> Tuple[np.ndarray, np.ndarray]:
        left_idxs = np.argwhere(X_column <= threshold).flatten()
        right_idxs = np.argwhere(X_column > threshold).flatten()
        return left_idxs, right_idxs
    def _entropy(self, y: np.ndarray) -> float:
        hist = np.bincount(y)
        ps = hist / len(y)
        return -np.sum([p * np.log2(p) for p in ps if p > 0])
    def _most_common_label(self, y: np.ndarray) -> Any:
        counter = Counter(y)
        return counter.most_common(1)[0][0]
    def predict(self, X: np.ndarray) -> np.ndarray:
        return np.array([self._traverse_tree(x, self.tree) for x in X])
    def _traverse_tree(self, x: np.ndarray, node: Node) -> Any:
        if node.value is not None:
            return node.value
        if x[node.feature] <= node.threshold:
            return self._traverse_tree(x, node.left)
        return self._traverse_tree(x, node.right)
# 示例数据集
X = np.array([
    [1, 1], [1, 1], [1, 0], [0, 1], [0, 1]
])
y = np.array(['Y', 'Y', 'N', 'N', 'N'])
# 训练模型
dt = DecisionTreeID3(max_depth=3)
dt.fit(X, y)
# 测试数据
X_test = np.array([[1, 0], [0, 0],[1,2]])
# 预测
predictions = dt.predict(X_test)
for i, prediction in enumerate(predictions):
    print(f"测试数据点 {X_test[i]} 的预测结果是 {prediction}")
